#!/usr/bin/python
# -*- coding: UTF-8 -*-

from PIL import Image, ImageDraw
import os
import random
import math
import multiprocessing
import xml.dom.minidom

def augmentation_flip(inputfolder, outputfolder, tmpfolder, filename):
    '''
    对指定目录下的指定名字的jpg和xml文件进行翻转形式的数据增广，并保存到指定目录下
    :param inputfolder:
    :param outputfolder:
    :param filename:
    :return:
    '''
    #打开图片
    inputfilepath_jpg = inputfolder + '/' + filename + '.jpg'
    im = Image.open(inputfilepath_jpg)

    #打开标注数据xml文件
    inputfilepath_xml = inputfolder + '/' + filename + '.xml'
    DOMTree = xml.dom.minidom.parse(inputfilepath_xml)
    collection = DOMTree.documentElement
    width = collection.getElementsByTagName("width")
    height = collection.getElementsByTagName("height")
    xmlxmin = collection.getElementsByTagName("xmin")
    xmlxmax = collection.getElementsByTagName("xmax")
    xmlymin = collection.getElementsByTagName("ymin")
    xmlymax = collection.getElementsByTagName("ymax")
    rectNum = len(xmlxmin)
    w = float(width[0].childNodes[0].data)
    h = float(height[0].childNodes[0].data)

    #图片和数据进行增广：水平翻转
    outputfilepath_jpg = outputfolder + '/' + filename + '_A.jpg'
    tmpfilepath_jpg = tmpfolder + '/' + filename + '_A.jpg'
    outputfilepath_xml = outputfolder + '/' + filename + '_A.xml'
    DOMTreeClone = DOMTree.cloneNode(10)

    im_flip = im.transpose(Image.FLIP_LEFT_RIGHT)
    im_flip.save(outputfilepath_jpg, "JPEG")

    # 图像的标注数据同步变动保存
    DOMTreeClone.documentElement.getElementsByTagName("filename")[0].childNodes[0].data = filename + '_A.jpg'
    #DOMTreeClone.documentElement.getElementsByTagName("path")[0].childNodes[0].data = str(outputfilepath_jpg)
    DOMTreeClone.documentElement.getElementsByTagName("width")[0].childNodes[0].data = str(int(w))
    DOMTreeClone.documentElement.getElementsByTagName("height")[0].childNodes[0].data = str(int(h))
    for i in range(rectNum):
        DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data = str(int(w - 1 - float(xmlxmax[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data = str(int(w - 1 - float(xmlxmin[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data = str(int(float(xmlymin[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data = str(int(float(xmlymax[i].childNodes[0].data)))
    with open(outputfilepath_xml, 'w', encoding='utf-8') as f:
        DOMTreeClone.writexml(f, encoding='utf-8')

    # 画框用于测试效果
    '''draw = ImageDraw.Draw(im_flip)
    for i in range(rectNum):
        xl = float(DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data)
        xr = float(DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data)
        yu = float(DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data)
        yd = float(DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data)
        draw.line([(xl, yu), (xr, yu)], "red")
        draw.line([(xr, yu), (xr, yd)], "red")
        draw.line([(xr, yd), (xl, yd)], "red")
        draw.line([(xl, yd), (xl, yu)], "red")
        im_flip.save(tmpfilepath_jpg, "JPEG")'''


    # 图片和数据进行增广：垂直翻转
    outputfilepath_jpg = outputfolder + '/' + filename + '_B.jpg'
    tmpfilepath_jpg = tmpfolder + '/' + filename + '_B.jpg'
    outputfilepath_xml = outputfolder + '/' + filename + '_B.xml'
    DOMTreeClone = DOMTree.cloneNode(10)

    im_flip = im.transpose(Image.FLIP_TOP_BOTTOM)
    im_flip.save(outputfilepath_jpg, "JPEG")

    # 图像的标注数据同步变动保存
    DOMTreeClone.documentElement.getElementsByTagName("filename")[0].childNodes[0].data = filename + '_B.jpg'
    #DOMTreeClone.documentElement.getElementsByTagName("path")[0].childNodes[0].data = outputfilepath_jpg
    DOMTreeClone.documentElement.getElementsByTagName("width")[0].childNodes[0].data = str(int(w))
    DOMTreeClone.documentElement.getElementsByTagName("height")[0].childNodes[0].data = str(int(h))
    for i in range(rectNum):
        DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data = str(int(float(xmlxmin[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data = str(int(float(xmlxmax[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data = str(int(h - 1 - float(xmlymax[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data = str(int(h - 1 - float(xmlymin[i].childNodes[0].data)))
        #print("dsgfdgfg")
    with open(outputfilepath_xml, 'w', encoding='utf-8') as f:
        DOMTreeClone.writexml(f, encoding='utf-8')

    # 画框用于测试效果
    '''draw = ImageDraw.Draw(im_flip)
    for i in range(rectNum):
        xl = float(DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data)
        xr = float(DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data)
        yu = float(DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data)
        yd = float(DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data)
        draw.line([(xl, yu), (xr, yu)], "white")
        draw.line([(xr, yu), (xr, yd)], "white")
        draw.line([(xr, yd), (xl, yd)], "white")
        draw.line([(xl, yd), (xl, yu)], "white")
        im_flip.save(tmpfilepath_jpg, "JPEG")'''


    # 图片和数据进行增广：水平+垂直翻转
    outputfilepath_jpg = outputfolder + '/' + filename + '_C.jpg'
    tmpfilepath_jpg = tmpfolder + '/' + filename + '_C.jpg'
    outputfilepath_xml = outputfolder + '/' + filename + '_C.xml'
    DOMTreeClone = DOMTree.cloneNode(10)

    im_flip = im.transpose(Image.ROTATE_180)
    im_flip.save(outputfilepath_jpg, "JPEG")

    # 图像的标注数据同步变动保存
    DOMTreeClone.documentElement.getElementsByTagName("filename")[0].childNodes[0].data = filename + '_C.jpg'
    #DOMTreeClone.documentElement.getElementsByTagName("path")[0].childNodes[0].data = outputfilepath_jpg
    DOMTreeClone.documentElement.getElementsByTagName("width")[0].childNodes[0].data = str(int(w))
    DOMTreeClone.documentElement.getElementsByTagName("height")[0].childNodes[0].data = str(int(h))
    for i in range(rectNum):
        DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data = str(int(w - 1 - float(xmlxmax[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data = str(int(w - 1 - float(xmlxmin[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data = str(int(h - 1 - float(xmlymax[i].childNodes[0].data)))
        DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data = str(int(h - 1 - float(xmlymin[i].childNodes[0].data)))
    with open(outputfilepath_xml, 'w', encoding='utf-8') as f:
        DOMTreeClone.writexml(f, encoding='utf-8')

    # 画框用于测试效果
    '''draw = ImageDraw.Draw(im_flip)
    for i in range(rectNum):
        xl = float(DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data)
        xr = float(DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data)
        yu = float(DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data)
        yd = float(DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data)
        draw.line([(xl, yu), (xr, yu)], "green")
        draw.line([(xr, yu), (xr, yd)], "green")
        draw.line([(xr, yd), (xl, yd)], "green")
        draw.line([(xl, yd), (xl, yu)], "green")
        im_flip.save(tmpfilepath_jpg, "JPEG")'''





def findfiles(path, t):
    '''
    遍历查找目录下的所有jpg文件名
    :param path:
    :param t:
    :return:
    '''
    files = os.listdir(path)
    for f in files:
        npath = path + '/' + f
        if(os.path.isfile(npath)):
            if(os.path.splitext(npath)[1] ==".jpg"):
                t.append(os.path.splitext(f)[0])
                #print(os.path.splitext(f)[0])
    return

if __name__=="__main__":
    '''
    主函数
    输入：带标注信息的图像
    输出：图像翻转后的一组图像，以及同步得到的对应标注信息，放入outputfiles下        
    '''
    print("version:2018.11.23 auther:lhf")

    # 图像翻转增广后的一组图像，以及同步得到的对应标注信息，放入outputfiles下
    inputfolder = "./trainData"
    outputfolder = "./augFiles"
    tmpfolder = "./tmpFiles"

    filenamelist = []
    inputfolder_copy = inputfolder
    findfiles(inputfolder_copy, filenamelist)

    # 按照翻转的方式对原始图像及标注数据进行增广
    pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
    for filename in filenamelist:
        outputfolder_copy = outputfolder
        tmpfolder_copy = tmpfolder
        if (os.path.exists(outputfolder_copy)):
            x = 1
        else:
            os.makedirs(outputfolder_copy)
            #os.makedirs(tmpfolder_copy)

        #print("processing %s ..." % filename)
        pool.apply_async(augmentation_flip, (inputfolder_copy, outputfolder_copy, tmpfolder_copy, filename))
    pool.close()
    pool.join()
    filenamelist.clear()
